import os
import numpy as np
import scipy.io.wavfile as wavfile
import librosa
from scipy.signal import find_peaks

input_folder = r'D:\Lv\Downloads\StethoData 拉萨古荣小学001 (2)\StethoData 拉萨古荣小学001'
output_folder = r'D:\Lv\Desktop\心音分割\StethoData 拉萨古荣小学1_切割'
default_Fs = 11025  # 默认采样率
n = 2  # 默认心跳数量


def process_files(input_folder, output_folder, n=2, default_Fs=11025):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 获取指定文件夹内的所有pcm和wav文件
    files = [f for f in os.listdir(input_folder) if f.endswith('.pcm') or f.endswith('.wav')]

    for file in files:
        print(f'Processing file: {file}')
        try:
            file_path = os.path.join(input_folder, file)
            if file.endswith('.pcm'):
                # 读取pcm文件
                with open(file_path, 'rb') as fid:
                    audio = np.frombuffer(fid.read(), dtype=np.int16)
                y0 = audio.astype(np.float32) / (2 ** 15)
                Fs = default_Fs  # 手动设置采样率
            else:
                # 读取wav文件
                y0, Fs = librosa.load(file_path, sr=None)

            # 找到心跳峰值，并获取它们在信号中的位置
            peaks, _ = find_peaks(y0, distance=int(0.15 * Fs), height=0.25 * np.max(y0))

            # 在每n次心跳之间切割信号，并保存为新的wav文件
            for i in range(1, len(peaks) - 2 * n + 1):
                # 计算开始和结束的样本点，向前和向后扩展0.15秒的样本数量
                start_sample = max(peaks[i] - int(0.15 * Fs), 0)
                end_sample = min(peaks[i + 2 * n - 1] + int(0.15 * Fs), len(y0))
                segment = y0[start_sample:end_sample]

                # 定义新的文件名和路径
                new_file_name = f'{os.path.splitext(file)[0]}_segment_{i}.wav'
                new_file_path = os.path.join(output_folder, new_file_name)

                # 保存切割的信号为新的wav文件
                wavfile.write(new_file_path, Fs, (segment * (2 ** 15)).astype(np.int16))

        except Exception as e:
            # 如果循环内的代码出现错误，跳过当前迭代并继续执行
            print(f'Error occurred while processing file: {file}. Skipping this file. Error: {e}')


# 调用函数处理文件
process_files(input_folder, output_folder, n, default_Fs)
